import pygame


class CustomGroup(pygame.sprite.Group):
    """
    在原有的基础上，进行了如下扩展：
    1. 按顺序存储 Sprite 的 Group，病新增了 pop() 方法。
    2. 重写 update 方法，使其返回调用了 update 方法后的 Sprite 对象列表。
    """

    def __init__(self):
        self.sprites_list = []
        super().__init__()

    def add(self, *sprites):
        super().add(*sprites)
        self.sprites_list.extend(sprites)

    def remove(self, *sprites):
        super().remove(*sprites)
        for sprite in sprites:
            if sprite in self.sprites_list:
                self.sprites_list.remove(sprite)

    def pop(self):
        if self.sprites_list:
            last_sprite = self.sprites_list.pop()
            self.remove(last_sprite)
            return last_sprite
        return None

    def update(self, *args, **kwargs):
        temp_list = []
        for sprite in self.sprites():
            temp_list.append(sprite.update(*args, **kwargs))
        return temp_list


class AnimationSprite(pygame.sprite.Sprite):
    """
    用于播放动画的精灵。
    """

    def __init__(self, images, pos=(0, 0), frame_duration=50, loop_count=0):
        """
        :param images: 帧图片列表（按顺序）
        :param pos: 动画绘制位置（中间坐标）
        :param frame_duration: 每帧持续的时间，单位：毫秒
        :param loop_count: 控制是否循环播放。-1 表示一直循环，0 表示不循环（只播放一次），1 表示再播放一次（只播放两次）。以此类推
        """
        super().__init__()

        # 所有帧图片
        self.images = images
        # 每帧持续时间
        self.frame_duration = frame_duration
        # 是否循环
        self.loop_count = loop_count
        # 当前图片帧
        self.current_frame = 0
        # 上一帧播放的时间
        self.last_animation_time = 0
        self.image = self.images[0]
        self.rect = self.image.get_rect()
        self.rect.center = pos
        # 动画是否正在播放
        self.is_running = False
        # 非循环时是否播放结束
        self.is_finished = False
        # 是否能被删除 (有时需要在动画播放完后，在被删除之前执行一些操作)
        self.can_be_deleted = False

    def start(self):
        """重新开始动画"""
        self.current_frame = 0
        self.last_animation_time = pygame.time.get_ticks()
        self.is_finished = False
        self.is_running = True

    def stop(self):
        """暂停动画"""
        self.is_running = False

    def update(self):
        """更新动画状态（每帧调用，自动处理帧切换）"""
        # 暂停或播放结束时不更新
        if not self.is_running or self.is_finished:
            return

        # 达到帧持续时间，切换到下一帧
        current_time = pygame.time.get_ticks()
        if current_time - self.last_animation_time >= self.frame_duration:
            self.current_frame += 1
            self.last_animation_time = current_time

            # 检查是否播放结束
            if self.current_frame >= len(self.images):
                if self.loop_count == -1:
                    self.current_frame = 0
                elif self.loop_count:
                    self.current_frame = 0
                    self.loop_count -= 1
                else:
                    self.current_frame = len(self.images) - 1
                    self.is_finished = True
                    self.is_running = False

        # 更新当前显示的图像
        self.image = self.images[self.current_frame]

    def draw(self, surface):
        """绘制动画到指定表面"""
        surface.blit(self.image, self.rect)
